import torch
import torch.nn as nn
from torch_scatter import scatter_add, scatter_mean, scatter_max,scatter_min
from utils.utilities import (
    decompose_and_trans_node_attr_to_cell_attr_graph,
    calc_cell_centered_with_node_attr,
)
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv

class GraphSCA3D(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()

        self.channel_excitation = nn.Sequential(
            nn.Linear(channel, int(channel // reduction)),
            nn.ReLU(inplace=True),
            nn.Linear(int(channel // reduction), channel),
        )
        self.spatial_se = GCNConv(
                    in_channels=channel,
                    out_channels=1,
        )

    def forward(self, x, batch, edge_index):
        BN,C = x.size()
        chn_se = scatter_mean(x,index=batch,dim=0).view(-1, C)
        chn_se = torch.sigmoid(self.channel_excitation(chn_se))
        chn_se = x*chn_se[batch]
        spa_se = torch.sigmoid(self.spatial_se(x,edge_index))
        spa_se = x*spa_se
        net_out = spa_se + x + chn_se
        return net_out
    
class GraphCA3D(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()

        self.channel_excitation = nn.Sequential(
            nn.Linear(channel, int(channel // reduction)),
            nn.ReLU(inplace=True),
            nn.Linear(int(channel // reduction), channel),
        )

    def forward(self, x, batch):
        BN,C = x.size()
        chn_se = scatter_mean(x,index=batch,dim=0).view(-1, C)
        chn_se = torch.sigmoid(self.channel_excitation(chn_se))
        chn_se = x*chn_se[batch]
        net_out = x + chn_se
        return net_out
    
class NodeBlock(nn.Module):
    def __init__(
        self, input_size, attention_size, attention=True, MultiHead=1, custom_func=None
    ):
        super(NodeBlock, self).__init__()
        self.net = custom_func
        # self.proj = nn.Linear(input_size, 3*input_size)
    def forward(self, graph_node, graph_cell=None):
    
        # Decompose graph
        (
            node_attr,
            edge_index,
            edge_attr,
            face,
            _,
            _,
        ) = decompose_and_trans_node_attr_to_cell_attr_graph(
            graph_node, has_changed_node_attr_to_cell_attr=True
        )

        ''' cell-based two step message passing algorithm '''
        if graph_cell is not None:
            senders_cell_idx, receivers_cell_idx = graph_cell.edge_index
            twoway_edge_attr = torch.cat((torch.chunk(edge_attr, 2, dim=-1)), dim=0)

            twoway_cell_connections_indegree = torch.cat(
                [senders_cell_idx, receivers_cell_idx], dim=0
            )

            twoway_cell_connections_outdegree = torch.cat(
                [receivers_cell_idx, senders_cell_idx], dim=0
            )

            cell_agg_received_edges = scatter_add(
                twoway_edge_attr, twoway_cell_connections_indegree, dim=0
            )

            cell_agg_neighbour_cell = scatter_add(
                cell_agg_received_edges[twoway_cell_connections_indegree],
                twoway_cell_connections_outdegree,
                dim=0,
            )

            cells_node = graph_node.face[0]
            cells_index = graph_cell.face[0]
            cell_to_node = cell_agg_neighbour_cell[cells_index]
            node_agg_received_edges = scatter_mean(cell_to_node, cells_node, dim=0)
            
            x = self.net(torch.cat((node_agg_received_edges, node_attr), dim=1))

        else:
            ''' node-based two step message passing algorithm'''

            senders_node_idx, receivers_node_idx = edge_index
            twoway_node_connections_indegree = torch.cat(
                [senders_node_idx, receivers_node_idx], dim=0
            )

            twoway_node_connections_outdegree = torch.cat(
                [receivers_node_idx, senders_node_idx], dim=0
            )
                  
            # sum agg
            twoway_edge_attr = torch.cat((torch.chunk(edge_attr, 2, dim=-1)), dim=0)
            node_agg_received_edges = scatter_add(
                twoway_edge_attr, 
                twoway_node_connections_indegree, 
                dim=0,
                out=torch.zeros((node_attr.size(0),twoway_edge_attr.size(1)),device=node_attr.device)
            )

            node_avg_neighbour_node = scatter_mean(
                node_agg_received_edges[twoway_node_connections_outdegree],
                twoway_node_connections_indegree,
                dim=0,
                out=torch.zeros((node_attr.size(0),twoway_edge_attr.size(1)),device=node_attr.device)
            )
            
            # update node attr
            x = self.net(torch.cat((node_avg_neighbour_node, 
                                    node_attr), 
                                   dim=1))
            
        return Data(
            x=x,
            edge_attr=edge_attr,
            edge_index=edge_index,
            face=face,
            num_graphs=graph_node.num_graphs,
            batch=graph_node.batch,
        )

class EdgeBlock(nn.Module):
    def __init__(self, input_size=None, custom_func=None):
        super(EdgeBlock, self).__init__()
        self.net = custom_func

    def forward(self, graph_node, graph_cell=None):
        (
            node_attr,
            edge_index,
            edge_attr,
            face,
            _,
            _,
        ) = decompose_and_trans_node_attr_to_cell_attr_graph(
            graph_node, has_changed_node_attr_to_cell_attr=True
        )

        edges_to_collect = []

        """ >>> node to cell and concancentendate to edge >>> """
        if graph_cell is not None:
        # aggrate node attr to cell
            cells_node = graph_node.face
            cells_index = graph_cell.face
            cell_attr = calc_cell_centered_with_node_attr(
                node_attr=node_attr,
                cells_node=cells_node,
                cells_index=cells_index,
                reduce="sum",
                map=True,
            )

            # concancentendate cell attr to edge
            senders_cell_idx, receivers_cell_idx = graph_cell.edge_index
            # filter self-loop face
            mask = torch.logical_not(senders_cell_idx == receivers_cell_idx).unsqueeze(1)

            senders_attr = cell_attr[senders_cell_idx]
            receivers_attr = cell_attr[receivers_cell_idx]

            edges_to_collect.append(senders_attr)
            edges_to_collect.append(receivers_attr * (mask.long()))
            edges_to_collect.append(edge_attr)
            """ <<< node to cell and concancentendate to edge <<< """
            
            collected_edges = torch.cat(edges_to_collect, dim=1)

            edge_attr_ = self.net(collected_edges)  # Update
            
        else:
            """ >>> only node concancentendate to edge >>> """
            senders_node_idx,receivers_node_idx= edge_index
            
            twoway_node_connections_indegree = torch.cat(
                [senders_node_idx, receivers_node_idx], dim=0
            )

            twoway_node_connections_outdegree = torch.cat(
                [receivers_node_idx, senders_node_idx], dim=0
            )
            
            node_avg_neighbour_node = scatter_add(
                node_attr[twoway_node_connections_outdegree],
                twoway_node_connections_indegree,
                dim=0,
                out=torch.zeros((node_attr.size(0),node_attr.size(1)),device=node_attr.device),
            )

            senders_attr = node_avg_neighbour_node[senders_node_idx]
            receivers_attr = node_avg_neighbour_node[receivers_node_idx]

            edges_to_collect.append(senders_attr)
            edges_to_collect.append(receivers_attr)
            edges_to_collect.append(edge_attr)
            """ >>>> only node concancentendate to edge >>> """

            collected_edges = torch.cat(edges_to_collect, dim=1)

            edge_attr_ = self.net(collected_edges)  # Update
            
        return Data(
            x=node_attr,
            edge_attr=edge_attr_,
            edge_index=edge_index,
            face=face,
            num_graphs=graph_node.num_graphs,
            batch=graph_node.batch,
        )
